#include "signature.h"
#include "smalloc.h"
#include "trace.h"

#include <sys/types.h>
#include <dirent.h>
#include <cstring>
#include <cstdio>

security_manager::security_manager(const std::string &privkey_file,
        const std::string &passphrase, const std::string &pubkey_dir)
    : privkey_(NULL)
{
    OpenSSL_add_all_algorithms();

    FILE *fp = NULL;
    char *pstr = strdup(passphrase.c_str());

    TDEBUG("initializing security manager: privkey: %s, pubkey_dir: %s",
            privkey_file.c_str(), pubkey_dir.c_str());

    // load the private key
    if ((fp = fopen(privkey_file.c_str(), "r")) == NULL) {
        throw std::string("failed to open private key file");
    }
    privkey_ = PEM_read_PrivateKey(fp, NULL, NULL, pstr);
    free(pstr);
    fclose(fp);
    if (privkey_ == NULL) {
        throw std::string("perhaps invalid pem file");
    }

    // load public keys
    DIR *dp;
    struct dirent *dirp;
    EVP_PKEY *k;
    if ((dp = opendir(pubkey_dir.c_str())) == NULL) {
        throw std::string("bad public key directory: " + pubkey_dir);
    }
    while ((dirp = readdir(dp)) != NULL) {
        std::string fname(dirp->d_name);
        if (fname == "." || fname == "..")
            continue;
        size_t sep = fname.rfind(".pub.pem");
        if (sep != std::string::npos) {
            k = load_pubkey(pubkey_dir+ "/" + fname);
            if (k != NULL) {
                std::string id = fname.substr(0, sep);
                TINFO("loaded public key for %s", id.c_str());
                pubkeys_[id] = k;
            } else {
                TERR("invalid public key file: %s", fname.c_str());
            }
        } else {
            TWARN("non-pem file %s ignored", fname.c_str());
        }
    }
    if (closedir(dp) == -1) {
        TWARN("could not properly close opened dir: %s", pubkey_dir.c_str());
    }
}

security_manager::~security_manager()
{
    EVP_cleanup();
    if (privkey_) EVP_PKEY_free(privkey_);
    using namespace std;
    for (map<string, EVP_PKEY *>::const_iterator itr = pubkeys_.begin(),
            iend = pubkeys_.end(); itr != iend; ++itr) {
        EVP_PKEY_free(itr->second);
    }
    pubkeys_.clear();
}

bool security_manager::sign(const void *buf, unsigned buf_sz, void **sig,
        unsigned *sig_sz)
{
    EVP_MD_CTX ctx;
    bool result = false;

    *sig = smalloc(EVP_PKEY_size(privkey_));
    EVP_MD_CTX_init(&ctx);
    if (EVP_SignInit(&ctx, EVP_sha1()) &&
            EVP_SignUpdate(&ctx, buf, buf_sz) &&
            EVP_SignFinal(&ctx, (unsigned char *)*sig, sig_sz, privkey_))
    {
        result = true;
    }
    EVP_MD_CTX_cleanup(&ctx);
    if (!result) {
        free(*sig);
        *sig = NULL;
        *sig_sz = 0;
    }

    return result;
}

bool security_manager::verify(const void *buf, unsigned buf_sz, void *sig,
        unsigned sig_sz, const char *sender)
{
    if (pubkeys_.find(sender) == pubkeys_.end())
        return false;

    EVP_PKEY *pubkey = pubkeys_[sender];
    bool result = false;
    EVP_MD_CTX ctx;
    EVP_MD_CTX_init(&ctx);
    result = (EVP_VerifyInit(&ctx, EVP_sha1()) &&
            EVP_VerifyUpdate(&ctx, buf, buf_sz) &&
            EVP_VerifyFinal(&ctx, (unsigned char *)sig, sig_sz, pubkey));
    EVP_MD_CTX_cleanup(&ctx);

    return result;
}

EVP_PKEY *security_manager::load_pubkey(const std::string &file)
{
    EVP_PKEY *pubkey = NULL;

    char *nopwd = strdup("");
    FILE *fp;

    if ((fp = fopen(file.c_str(), "r")) == NULL) {
        TERR("cannot open pubkey file: %s", file.c_str());
    } else {
        pubkey = PEM_read_PUBKEY(fp, NULL, NULL, nopwd);
        fclose(fp);
    }
    free(nopwd);

    return pubkey;
}

